#!/usr/bin/python
# plot_fitc0rho0.py
'''Plots the outcome of GM_fitc0rho0.'''

import numpy
from matplotlib import use as useBackend
useBackend('Agg')
import matplotlib.pyplot as ppt


def cheat111():
    '''Cheat to clear the plot.'''
    ppt.subplot(311)
    ppt.plot([0,1],[1,2])
    ppt.subplot(131)
    ppt.plot([0,1],[1,2])
    ppt.subplot(111)

def choosecolor(j, colorlist):
    ncolors = len(colorlist)
    while True:
        if j > (ncolors - 1):
            j = j - ncolors
        else:
            colornow = colorlist[j]
            break
    return colornow

def plot(simpars, phypars, difflist, distlist, distarraylist, paramlist, clocslist, rlocslist):
    # Intake etc.
    for i in simpars:
        cmd = "%s = simpars['%s']" % (i,i)
        exec cmd
    for i in phypars:
        cmd = "%s = phypars['%s']" % (i,i)
        exec cmd
    colorlist = ['k', 'b', 'g', 'r', 'c', 'm', 'y', 'indigo', 'darkorange', 'saddlebrown', 'forestgreen', 'deeppink', 'slategray', 'maroon', 'lime'] #15
    difflist = numpy.array(difflist)
    distlist = numpy.array(distlist)
    distarraylist = numpy.array(distarraylist)
    paramlist = numpy.array(paramlist)
    length = len(difflist)
    xlist = numpy.arange(0, length+1)

    # Plot evolution of fit
    cheat111()
    ppt.subplot(211)
    ppt.title('Evolution of c0 during fitting routine')
    ppt.plot(xlist, paramlist[:,0])
    ppt.ylabel('Best c0')
    ppt.subplot(212)
    ppt.plot(xlist[:-1], numpy.log10(paramlist[:-1,2]))
    ppt.ylabel('$\log_{10}$ c0 range tested')
    ppt.xlabel('Iteration number')
    plotname = outdir+'FitEvolution_c0_'+timestamp+'.pdf'
    ppt.savefig(plotname)

    cheat111()
    ppt.subplot(211)
    ppt.title('Evolution of rho0 during fitting routine')
    ppt.plot(xlist, paramlist[:,1])
    ppt.ylabel('Best rho0')
    ppt.subplot(212)
    ppt.plot(xlist[:-1], numpy.log10(paramlist[:-1,3]))
    ppt.ylabel('$\log_{10}$ rho0 range tested')
    ppt.xlabel('Iteration number')
    plotname = outdir+'FitEvolution_rho_'+timestamp+'.pdf'
    ppt.savefig(plotname)

    cheat111()
    ppt.title('Evolution of goodness-of-fit during fitting routine')
    ppt.subplot(211)
    ppt.plot(xlist[1:], numpy.log10(distlist))
    ppt.ylabel('''$\log_{10}$ "distance"
from data''')
    ppt.subplot(212)
    ppt.plot(xlist[1:], numpy.log10(difflist))
    ppt.ylabel('''$\log_{10}$ difference between
two best test points''')
    ppt.xlabel('Iteration number')
    plotname = outdir+'FitEvolution_dist_'+timestamp+'.pdf'
    ppt.savefig(plotname)

    # Plot how location jumps around
    cheat111()
    locslength = len(clocslist)
    count = len(clocslist[0,:])
    plotlist = []
    plotnames = []
    for i in range(0, locslength):
        c = numpy.log10(clocslist[i])
        r = numpy.log10(rlocslist[i])
        colornow = choosecolor(i, colorlist)
        if i < 2:
            mkr = '.'
        else:
            mkr = ','
        for j in range(0, count):
            tc = c[j]*numpy.ones(count,dtype=float)
            lastplot = ppt.plot(r, tc, c=colornow, marker=mkr, ls='None')
        plotlist.append(lastplot)
        name = "Itrn. "+str(i+1)
        plotnames.append(name)
# For some reason this block of code doesn't appear to do... anything, so I have it commented out. I wish it worked, though.
#    for i in range(0, locslength):
#        colornow = choosecolor(i, colorlist)
#        c0i = paramlist[i,0]
#        rho0i = paramlist[i,1]
#        ppt.plot(rho0i, c0i, c=colornow, marker='x', ls='none')
    ppt.title('Locations sampled, by iteration')
    ppt.xlabel('$\log_{10}$ rho0')
    ppt.ylabel('$\log_{10}$ c0')
    ppt.legend(plotlist, plotnames)
    plotname = outdir+'FitLocations_'+timestamp+'.pdf'
    ppt.savefig(plotname)

    # Plot weighting scheme
    cheat111()
    intl = latitudeweightinternal/latitudeweightinternal.mean()
    extl = latitudeweightexternal/latitudeweightexternal.mean()
    totl1 = latitudeweightinternal*latitudeweightexternal
    totl = totl1/totl1.mean()
    internal = ppt.plot(bvals, intl, c='b', ls='-.', marker='o')
    external = ppt.plot(bvals, extl, c='g', ls='-.', marker='o')
    total = ppt.plot(bvals, totl, c='k', ls='-.', marker='o')
    plotlist = [internal,external,total]
    plotnames = ["Intrinsic weights", "Imposed weights", "Total weights used"]
    ppt.legend(plotlist, plotnames, loc=2)
    ppt.xlabel('Galactic latitude, degrees')
    ppt.ylabel('Relative weight assigned')
    titl = "Fitting weights for 'latitudeweight' value "+str(latitudeweight)
    ppt.title(titl)
    plotname = outdir+'FitWeights_'+timestamp+'.pdf'
    ppt.savefig(plotname)

    # Plot individual fit iterations
    for i in range(0, length):
        cheat111()
        ppt.imshow(distarraylist[i,:,:])
        ppt.colorbar()
        plottitle = "'Distance' surface for iteration "+str(i+1)
        ppt.title(plottitle)
        ppt.xlabel('Small -> rho0 -> Large')
        ppt.ylabel('Large -> c0 -> Small')
        plotname = outdir+'FitSurface_'+timestamp+'_Iteration'+str(i+1)+'.pdf'
        ppt.savefig(plotname)
